"""
keras model fit
"""
from keras.callbacks import TerminateOnNaN, EarlyStopping
from utils.plot_graph import plot_model_graph


def train_model(model, train_x, train_y, test_x, test_y, epochs, model_name,
                pic_path, shuffle=True, verbose=1, plot=False,
                batch_size=72, workers=8, use_multiprocessing=True):
    """
    keras model fit
    :param model:
    :param train_x:
    :param train_y:
    :param test_x:
    :param test_y:
    :param epochs:
    :param model_name:
    :param pic_path:
    :param shuffle:
    :param verbose:
    :param plot:
    :param batch_size:
    :param workers:
    :param use_multiprocessing:
    :return:
    """
    # add callback
    callback_list = [TerminateOnNaN(),
                     EarlyStopping(monitor='val_loss', patience=3, mode='auto')]
    history = model.fit(train_x, train_y,
                        epochs=epochs,
                        batch_size=batch_size,
                        validation_data=(test_x, test_y),
                        verbose=verbose,
                        workers=workers,
                        shuffle=shuffle,
                        use_multiprocessing=use_multiprocessing,
                        callbacks=callback_list)

    if plot:
        plot_model_graph(model_fit=history,
                         model_name=model_name,
                         pic_path=pic_path)

    loss_metrics = model.evaluate(test_x, test_y,
                                  verbose=0,
                                  batch_size=batch_size,
                                  workers=workers,
                                  use_multiprocessing=use_multiprocessing)

    return model, loss_metrics
